package at.tasat.assignment;

import java.util.Stack;

import at.tasat.cnf.Clause;
import at.tasat.cnf.Literal;
import at.tasat.util.IntKeyMap;
import at.tasat.util.IntStack;

public final class PartialAssignment implements Cloneable {

	private final int nVariables;
	private int decisionLevel;
	private final IntStack control;
	private final Stack<Assignment> trail;
	private final IntKeyMap<Assignment> assignments;

	public PartialAssignment(int nVariables) {
		this(nVariables, 0, new IntStack(nVariables + 1), new Stack<Assignment>());
		assert nVariables >= 0;
		control.push(0);
	}

	private PartialAssignment(int nVariables, int decisionLevel, IntStack control, Stack<Assignment> trail) {
		this.nVariables = nVariables;
		this.decisionLevel = decisionLevel;
		this.control = control;
		this.trail = trail;
		assignments = new IntKeyMap<Assignment>(1, nVariables);
		for (Assignment assignment : trail) {
			assignments.put(assignment.getVariableId(), assignment);
		}
	}

	private boolean isSane() {
		return decisionLevel >= 0 && control.size() == decisionLevel + 1 && trail.size() >= control.peek();
	}

	public int getDecisionLevel() {
		return decisionLevel;
	}

	public Assignment getDecision(int index) {
		assert index > 0;
		return trail.get(control.get(index));
	}

	public int getNumberOfCurrentDecisionAssignments() {
		return trail.size() - control.peek();
	}

	public boolean isAssigned(int variableId) {
		return assignments.containsKey(variableId);
	}

	public Assignment getAssignmentOfVariable(int variableId) {
		return assignments.get(variableId);
	}

	public int getNumberOfAssignments() {
		return trail.size();
	}

	public Assignment getAssignment(int index) {
		return trail.get(index);
	}

	public Assignment[] getAllAssignments() {
		return trail.toArray(new Assignment[trail.size()]);
	}

	public boolean isSatisfied(Clause clause) {
		for (int i = 0; i < clause.getNumberOfLiterals(); i++) {
			int literal = clause.getLiteral(i);
			int variableId = Literal.getVariableId(literal);
			if (assignments.containsKey(variableId) && assignments.get(variableId).isTrue(literal)) {
				return true;
			}
		}
		return false;
	}

	public boolean isUnsatisfied(Clause clause) {
		for (int i = 0; i < clause.getNumberOfLiterals(); i++) {
			int literal = clause.getLiteral(i);
			int variableId = Literal.getVariableId(literal);
			if (!assignments.containsKey(variableId) || assignments.get(variableId).isTrue(literal)) {
				return false;
			}
		}
		return true;
	}

	public int getUnitVariableId(Clause clause) {
		int unassignedVariableId = 0;
		for (int i = 0; i < clause.getNumberOfLiterals(); i++) {
			int literal = clause.getLiteral(i);
			int variableId = Literal.getVariableId(literal);
			if (assignments.containsKey(variableId)) {
				if (assignments.get(variableId).isTrue(literal)) {
					return 0;
				}
			} else {
				if (unassignedVariableId > 0) {
					return 0;
				} else {
					unassignedVariableId = variableId;
				}
			}
		}
		return unassignedVariableId;
	}

	public int getNumberOfUnassignedLiterals(Clause clause) {
		int unassigned = 0;
		for (int i = 0; i < clause.getNumberOfLiterals(); i++) {
			int variableId = Literal.getVariableId(clause.getLiteral(i));
			if (!isAssigned(variableId)) {
				unassigned++;
			}
		}
		return unassigned;
	}

	public Assignment pushAssignment(int variableId, boolean phase, boolean isDecision) {
		assert isSane();
		Assignment assignment = new Assignment(variableId, phase);
		assert !assignments.containsKey(assignment.getVariableId());
		assignments.put(variableId, assignment);
		trail.push(assignment);
		if (isDecision) {
			control.push(trail.size() - 1);
			decisionLevel++;
		}
		assert isSane();
		return assignment;
	}

	public int popAssignment() {
		assert isSane();
		Assignment assignment = trail.pop();
		assignments.remove(assignment.getVariableId());
		if (control.size() > 1 && control.peek() >= trail.size()) {
			control.pop();
			decisionLevel--;
		}
		assert isSane();
		return assignment.getVariableId();
	}
	
	public void copyTo(PartialAssignment to) {
		assert isSane();
		assert getNumberOfAssignments() > to.getNumberOfAssignments();
		int controlIndex = 1;
		for(int i = 0; i < trail.size(); i++){
			boolean decision = controlIndex < control.size() ? control.get(controlIndex) == i : false;
			if(decision){
				controlIndex++;
			}
			Assignment assignment = trail.get(i);
			if(i >= to.getNumberOfAssignments()){
				to.pushAssignment(assignment.getVariableId(), assignment.getPhase(), decision);
			} else {
				assert assignment.equals(to.getAssignment(i));
			}
		}
		assert isSane();
	}

	@Override
	@SuppressWarnings("unchecked")
	public PartialAssignment clone() {
		return new PartialAssignment(nVariables, decisionLevel, control.clone(), (Stack<Assignment>) trail.clone());
	}

	@Override
	public String toString() {
		return "decisionlevel = " + decisionLevel + ", control stack = " + toStringControl() + ", trail = " + toStringTrail();
	}

	private String toStringControl() {
		StringBuilder string = new StringBuilder();
		boolean first = true;
		for (int c : control) {
			if (first) {
				first = false;
			} else {
				string.append(", ");
			}
			string.append(c);
		}
		return '(' + string.toString() + ')';
	}

	private String toStringTrail() {
		StringBuilder string = new StringBuilder();
		boolean first = true;
		for (Assignment assignment : trail) {
			if (first) {
				first = false;
			} else {
				string.append(", ");
			}
			string.append(assignment);
		}
		return '(' + string.toString() + ')';
	}

}
